# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment3/question2/datasets.py
import torch.utils.data


class EpsDataset(torch.utils.data.Dataset):
    def __init__(self):
        self._data: list[tuple[torch.Tensor, torch.Tensor]] = []
        with open("./data/Eps_real.txt") as file_r:
            # noinspection SpellCheckingInspection
            with open("./data/Eps_imag.txt") as file_i:
                line = file_r.readline().strip()
                assert line == "#"
                line = file_r.readline().strip()
                assert line == '''#"Frequency / GHz"	"Eps' (Data list) [Real Part]"'''
                line = file_r.readline().strip()
                assert line == "#------------------------------------------------"

                line = file_i.readline().strip()
                assert line == "#"
                line = file_i.readline().strip()
                # 这里给的表头还是 Real Part
                assert line == '''#"Frequency / GHz"	"Eps'' (Data list) [Real Part]"'''
                line = file_i.readline().strip()
                assert line == "#-------------------------------------------------"

                line_r: str
                line_i: str
                for line_r, line_i in zip(file_r, file_i):
                    line_r_split = line_r.strip().split()
                    line_i_split = line_i.strip().split()
                    assert len(line_r_split) == len(line_i_split) == 2
                    assert line_r_split[0] == line_i_split[0]
                    x = float(line_r_split[0])
                    r = float(line_r_split[1])
                    i = float(line_i_split[1])
                    self._data.append((torch.tensor(x), torch.tensor([r, i])))

                assert file_r.readline().strip() == ""
                assert file_i.readline().strip() == ""

                self._data.pop(0)
                self._data.pop(1)
                self._data.pop(2)
                self._data.pop(3)
                self._data.pop(4)
                self._data.pop(5)

    def __getitem__(self, item):
        return self._data[item]

    def __len__(self):
        return len(self._data)
